%load_ext autoreload
%autoreload 2

import numpy
import torch
from adaptive_nof1 import *
from adaptive_nof1.policies import *
from adaptive_nof1.helpers import *
from adaptive_nof1.inference import *
from adaptive_nof1.metrics import *
from matplotlib import pyplot as plt
import seaborn

WARNING (pytensor.tensor.blas): Using NumPy C-API based implementation for BLAS functions.

# Setup generic n-of-1 parameters
block_length = 5
max_length = 10 * block_length
number_of_actions = 2
number_of_patients = 100

# Scenarios
class NormalModel(Model):
    def __init__(self, patient_id, mean, variance):
        self.rng = numpy.random.default_rng(patient_id)
        self.mean = mean
        self.variance = variance
        self.patient_id = patient_id

    def multivariate_normal_distribution(debug_data):
        cov = torch.diag_embed(torch.tensor(self.variance))
        return torch.distributions.MultivariateNormal(torch.tensor(self.mean), cov)

    def generate_context(self, history):
        return {}

    @property
    def additional_config(self):
        return {"expectations_of_interventions": self.mean}

    @property
    def number_of_interventions(self):
        return len(self.mean)

    def observe_outcome(self, action, context):
        treatment_index = action["treatment"]
        return {"outcome": self.rng.normal(self.mean[treatment_index], self.variance[treatment_index])}

    def __str__(self):
        return f"NormalModel({self.mean, self.variance})"

generating_scenario_I = lambda patient_id: NormalModel(patient_id, mean=[0, 0], variance=[1,1])
generating_scenario_II = lambda patient_id: NormalModel(patient_id, mean=[1, 0], variance=[1,1])
generating_scenario_III = lambda patient_id: NormalModel(patient_id, mean=[2, 0], variance=[1,1])

# Inference Model
inference_model = lambda: NormalKnownVariance(prior_mean=0, prior_variance=1, variance=1)

# Stopping Time
ALPHA_STOPPING = 0.01
def alpha_stopping_time(history, context):
    model = NormalKnownVariance(prior_mean=0, prior_variance=1, variance=1)
    model.update_posterior(history, number_of_actions)
    probabilities = model.approximate_max_probabilities(number_of_actions, context)
    return 1 - max(probabilities) < ALPHA_STOPPING

# Policies
fixed_policy = StoppingPolicy(
    policy = BlockPolicy(
        block_length = block_length,
        internal_policy = FixedPolicy(
            number_of_actions=2,
            inference_model = inference_model(),
        ),
    ),
    stopping_time = alpha_stopping_time,
)

explore_then_commit = StoppingPolicy(
    policy= BlockPolicy(
        block_length = block_length,
        internal_policy = ExploreThenCommit(
        number_of_actions=2,
        exploration_length=4,
        block_length = block_length,
        inference_model = inference_model(),
    
        ),
    ),
    stopping_time = alpha_stopping_time,
)


thompson_sampling_policy = StoppingPolicy(
        policy = BlockPolicy(
            block_length = block_length,
            internal_policy = ThompsonSampling(
                inference_model=inference_model(),
                number_of_actions=2,
            ),
        ),
    stopping_time = alpha_stopping_time,
)

ucb_policy = StoppingPolicy(
    policy = BlockPolicy(
        block_length = block_length,
        internal_policy = UpperConfidenceBound(
            inference_model=inference_model(),
            number_of_actions=2,
            epsilon=0.05,
        ),
    ),
    stopping_time = alpha_stopping_time,
)

# Full crossover study
study_designs = {
    "n_patients": [number_of_patients],
    "policy": [fixed_policy, explore_then_commit, thompson_sampling_policy, ucb_policy],
    "model_from_patient_id": [
        generating_scenario_I, generating_scenario_II, generating_scenario_III,
    ]
}
configurations = generate_configuration_cross_product(study_designs)

calculated_series, config_to_simulation_data = simulate_configurations(
    configurations, max_length
)

  0%|          | 0/50 [00:00<?, ?it/s]

  0%|          | 0/50 [00:00<?, ?it/s]

  0%|          | 0/50 [00:00<?, ?it/s]

  0%|          | 0/50 [00:00<?, ?it/s]

  0%|          | 0/50 [00:00<?, ?it/s]

  0%|          | 0/50 [00:00<?, ?it/s]

  0%|          | 0/50 [00:00<?, ?it/s]

  0%|          | 0/50 [00:00<?, ?it/s]

  0%|          | 0/50 [00:00<?, ?it/s]

  0%|          | 0/50 [00:00<?, ?it/s]

# Todo: make the output table in a way that we chose the maximum index
def debug_data_to_torch_distribution(debug_data):
    mean = debug_data["mean"]
    # + the true variance of 1
    variance = numpy.array(debug_data["variance"]) + 1
    cov = torch.diag_embed(torch.tensor(variance))
    return torch.distributions.MultivariateNormal(torch.tensor(mean), cov)

def data_to_true_distribution(data):
    mean = data.additional_config["expectations_of_interventions"]
    cov = torch.eye(len(mean))
    return torch.distributions.MultivariateNormal(torch.tensor(mean), cov)


metrics = [
    SimpleRegretWithMean(),
    BestArmIdentification(),
    CumulativeRegret(),
    Length(),
    KLDivergence(data_to_true_distribution = data_to_true_distribution, debug_data_to_posterior_distribution=debug_data_to_torch_distribution),
]
model_mapping = {
    "NormalModel(([0, 0], [1, 1]))": "I",
    "NormalModel(([1, 0], [1, 1]))": "II",
    "NormalModel(([2, 0], [1, 1]))": "III",
}
policy_mapping = {
    "StoppingPolicy(BlockPolicy(FixedPolicy))": "Fixed",
    "StoppingPolicy(BlockPolicy(ThompsonSampling(NormalKnownVariance(0, 1, 1))))": "TS",
    "StoppingPolicy(BlockPolicy(UpperConfidenceBound(0.05 epsilon, NormalKnownVariance(0, 1, 1))))": "UCB",
    "StoppingPolicy(BlockPolicy(ExploreThenCommit(4,NormalKnownVariance(0, 1, 1))))": "ETC",
}

df = SeriesOfSimulationsData.score_data(
    [s["result"] for s in calculated_series], metrics, {"model": lambda x: model_mapping[x], "policy": lambda x: policy_mapping[x]}
)

df = df.reset_index(drop=True)
max_t_indices = df.groupby(["policy", "metric", "model", "patient_id"])["t"].idxmax()
filtered_df = df.iloc[max_t_indices]
filtered_df = filtered_df.reset_index(drop=True)
groupby_columns = ["model", "policy"]

pivoted_df = filtered_df.pivot(
    index=["model", "policy", "patient_id"],
    columns="metric",
    values="score",
)
table = pivoted_df.groupby(groupby_columns).agg(['mean', 'std'])

policy_ordering = ["Fixed", "ETC", "SH", "UCB", "TS"]

# Convert the 'policy' column in the MultiIndex to a Categorical type with the specified order
table = table.reset_index()
table['policy'] = pd.Categorical(table['policy'], categories=policy_ordering, ordered=True)

# Sort the DataFrame first by 'model' then by the now-ordered 'policy'
sorted_table = table.sort_values(by=['model', 'policy']).set_index(groupby_columns)[["Cumulative Regret (outcome)", "KL Divergence", "Simple Regret With Mean", "Length", "Best Arm Identification With Mean"]].rename(
    columns={"Cumulative Regret (outcome)": "Regret", "Simple Regret With Mean": "$\SIR$", "Best Arm Identification With Mean": "$\BAI$", "KL Divergence": "$\KLD$"},
)
sorted_table.index.names = ["S.", "Policy"]
sorted_table

with open('mt_resources/7-stopping/01-table.tex', 'w') as file:
    str = sorted_table.style.format(precision=2).to_latex(hrules=True)
    print(str)
    file.write(str)

\begin{tabular}{lllrlrlrlrlr}
\toprule
 & metric & \multicolumn{2}{r}{Regret} & \multicolumn{2}{r}{$\KLD$} & \multicolumn{2}{r}{$\SIR$} & \multicolumn{2}{r}{Length} & \multicolumn{2}{r}{$\BAI$} \\
 &  & mean & std & mean & std & mean & std & mean & std & mean & std \\
S. & Policy &  &  &  &  &  &  &  &  &  &  \\
\midrule
\multirow[c]{4}{*}{I} & Fixed & -0.04 & 3.60 & 0.16 & 0.17 & 0.00 & 0.00 & 16.52 & 8.82 & 0.48 & 0.50 \\
 & ETC & -0.29 & 4.04 & 0.17 & 0.16 & 0.00 & 0.00 & 17.17 & 11.32 & 0.48 & 0.50 \\
 & UCB & 0.31 & 4.45 & 0.18 & 0.17 & 0.00 & 0.00 & 18.45 & 12.69 & 0.46 & 0.50 \\
 & TS & -0.15 & 4.33 & 0.15 & 0.16 & 0.00 & 0.00 & 22.42 & 14.84 & 0.52 & 0.50 \\
\multirow[c]{4}{*}{II} & Fixed & -6.73 & 4.48 & 0.22 & 0.13 & 0.04 & 0.20 & 11.08 & 4.86 & 0.96 & 0.20 \\
 & ETC & -6.91 & 4.51 & 0.22 & 0.14 & 0.05 & 0.22 & 11.27 & 4.92 & 0.95 & 0.22 \\
 & UCB & -6.50 & 5.64 & 0.40 & 0.28 & 0.14 & 0.35 & 11.60 & 6.54 & 0.86 & 0.35 \\
 & TS & -12.27 & 14.32 & 0.41 & 0.28 & 0.22 & 0.42 & 18.79 & 12.12 & 0.78 & 0.42 \\
\multirow[c]{4}{*}{III} & Fixed & -10.11 & 2.81 & 0.38 & 0.40 & 0.04 & 0.28 & 7.43 & 1.29 & 0.98 & 0.14 \\
 & ETC & -10.24 & 2.72 & 0.35 & 0.35 & 0.02 & 0.20 & 7.55 & 1.37 & 0.99 & 0.10 \\
 & UCB & -33.58 & 40.02 & 1.11 & 0.97 & 0.38 & 0.79 & 19.59 & 18.10 & 0.81 & 0.39 \\
 & TS & -41.30 & 42.31 & 1.16 & 1.02 & 0.50 & 0.87 & 25.07 & 18.01 & 0.76 & 0.43 \\
\bottomrule
\end{tabular}

def rename_df(df):
    df["policy_#_metric_#_model_p"] = df["policy"].apply(lambda x: policy_mapping[x])
    return df

SeriesOfSimulationsData.plot_lines(
    [s["result"] for s in calculated_series if s["configuration"]["model"] == "NormalModel(([1, 0], [1, 1]))"],
    [
        CumulativeRegret(),
    ],
    legend_position=(0.02,0.3),
    process_df = rename_df,
)
plt.ylabel('Regret')
plt.savefig("mt_resources/7-stopping/01_cumulative_regret.pdf", bbox_inches="tight")

SeriesOfSimulationsData.plot_lines(
    [s["result"] for s in calculated_series if s["configuration"]["model"] == "NormalModel(([1, 0], [1, 1]))"],
    [
        SimpleRegretWithMean(),
    ],
    legend_position=(0.8,1.0),
    process_df = rename_df,
)
plt.ylabel('Simple Regret')
plt.savefig("mt_resources/7-stopping/01_simple_regret.pdf", bbox_inches="tight")

SeriesOfSimulationsData.plot_lines(
    [s["result"] for s in calculated_series if s["configuration"]["model"] == "NormalModel(([1, 0], [1, 1]))"],
    [
        KLDivergence(data_to_true_distribution = data_to_true_distribution, debug_data_to_posterior_distribution=debug_data_to_torch_distribution)
    ],
    legend_position=(0.8,1.0),
    process_df = rename_df,
)
plt.ylabel('KL Divergence')
plt.savefig("mt_resources/7-stopping/01-kl-divergence.pdf", bbox_inches="tight")

df = SeriesOfSimulationsData.score_data(
    [s["result"] for s in calculated_series if s["configuration"]["model"] == "NormalModel(([1, 0], [1, 1]))"],
    [ IsStopped() ],
)
df["policy"] = df["policy"].apply(lambda x: policy_mapping[x])
groupby_df_sum = df.groupby(["policy", "model", "t"]).sum()

ax = seaborn.lineplot(
    data=groupby_df_sum,
    x="t",
    y="score",
    hue="policy",
    # units="patient_id",
    #estimator=numpy.median,
    #errorbar=lambda x: (numpy.quantile(x, 0.25), numpy.quantile(x, 0.75)),
)
plt.ylabel("Number of patients")
seaborn.move_legend(ax, "upper right", title=None)
plt.savefig("mt_resources/7-stopping/01_is_stopped.pdf", bbox_inches="tight")

plot_allocations_for_calculated_series(calculated_series)

/opt/homebrew/Caskroom/miniconda/base/envs/mt/lib/python3.11/site-packages/holoviews/plotting/bokeh/plot.py:987: UserWarning: found multiple competing values for 'toolbar.active_drag' property; using the latest value
  layout_plot = gridplot(
/opt/homebrew/Caskroom/miniconda/base/envs/mt/lib/python3.11/site-packages/holoviews/plotting/bokeh/plot.py:987: UserWarning: found multiple competing values for 'toolbar.active_scroll' property; using the latest value
  layout_plot = gridplot(

plot_allocations_for_calculated_series([s for s in calculated_series if s["configuration"]["policy"] == "StoppingPolicy(BlockPolicy(UpperConfidenceBound(0.05 epsilon, NormalKnownVariance(0, 1, 1))))" and s["configuration"]["model"] == "NormalModel(([0, 0], [1, 1]))"])

:Layout

plot_allocations_for_calculated_series([s for s in calculated_series if s["configuration"]["policy"] == "StoppingPolicy(BlockPolicy(UpperConfidenceBound(0.05 epsilon, NormalKnownVariance(0, 1, 1))))" and s["configuration"]["model"] == "NormalModel(([1, 0], [1, 1]))"])

:Layout

import param
from functools import reduce


class PatientExplorer(param.Parameterized):
    patient_id = param.Integer(default=0, bounds=(0, number_of_patients - 1))
    configuration = param.Integer(default=0, bounds=(0, len(calculated_series) - 1))
    t = param.Integer(default=10, bounds=(0, max_length - 1))

    @param.depends("patient_id", "configuration")
    def hvplot(self):
        debug_data = (
            calculated_series[self.configuration]["result"]
            .simulations[self.patient_id]
            .history.debug_data()
        )
        df = pandas.DataFrame([flatten_dictionary(d) for d in debug_data])
        return df.hvplot()

    @param.depends("configuration")
    def configuration_name(self):
        return panel.panel(calculated_series[self.configuration]["configuration"])

import holoviews
import panel
explorer = PatientExplorer()
hvplot = holoviews.DynamicMap(explorer.hvplot)
panel.Column(
    panel.Row(panel.Column(explorer.param, explorer.configuration_name), hvplot),
)

	metric	Regret		$\KLD$		$\SIR$		Length		$\BAI$
		mean	std	mean	std	mean	std	mean	std	mean	std
S.	Policy
I	Fixed	-0.037432	3.597176	0.163983	0.165737	0.0	0.000000	16.52	8.818575	0.48	0.502117
	ETC	-0.292172	4.041548	0.166418	0.164108	0.0	0.000000	17.17	11.315543	0.48	0.502117
	UCB	0.309196	4.454360	0.17932	0.170011	0.0	0.000000	18.45	12.693254	0.46	0.500908
	TS	-0.147157	4.328133	0.150667	0.159053	0.0	0.000000	22.42	14.843412	0.52	0.502117
II	Fixed	-6.732649	4.478590	0.219876	0.134231	0.04	0.196946	11.08	4.856902	0.96	0.196946
	ETC	-6.911719	4.506643	0.218885	0.135457	0.05	0.219043	11.27	4.921310	0.95	0.219043
	UCB	-6.502219	5.637445	0.403294	0.284822	0.14	0.348735	11.6	6.539700	0.86	0.348735
	TS	-12.266483	14.320192	0.411553	0.281109	0.22	0.416333	18.79	12.119185	0.78	0.416333
III	Fixed	-10.109516	2.809726	0.376703	0.399681	0.04	0.281411	7.43	1.289076	0.98	0.140705
	ETC	-10.243181	2.724291	0.354477	0.347126	0.02	0.200000	7.55	1.366075	0.99	0.100000
	UCB	-33.577203	40.023253	1.111306	0.969029	0.38	0.788554	19.59	18.100781	0.81	0.394277
	TS	-41.302591	42.305474	1.16179	1.023429	0.5	0.870388	25.07	18.014167	0.76	0.429235

Master's thesis case study 3: Bandit's with stopping¶